#include <iostream>

#include <vector>
#include <map>
#include <set>
#include <memory>

#include <cassert>

#include <sds/sfile.h>
#include <sds/sdataset.h>
#include <sds/sdataview.h>

#include <matrix.h>

#include <cluster.h>
#include <segment.h>

#include <lemon/grid_graph.h>
#include <lemon/unionfind.h>
#include <lemon/random.h>

static const double TH = 20.0;

#if !defined(DIST_EUCLID) && !defined(DIST_NORM)
#  define DIST_NORM 
#endif
using namespace lemon;
using namespace std;
using namespace sds;

typedef GridGraph Graph;
GRAPH_TYPEDEFS(Graph);

typedef Graph::NodeMap<int> IntNodeMap;
typedef Graph::NodeMap<Point> PointNodeMap;
typedef Graph::EdgeMap<double> WeightEdgeMap;

void Read(const string& filename, Graph& graph, 
	  WeightEdgeMap& weight, vector<SInteger>& size) {
  SFile file(filename);

  SDataView<unsigned char> source[CH];
  for (int k = 0; k < CH; ++k) {
    SDataset dataset = file.getDataset(k);
    source[k] = dataset.getDataView<unsigned char>();
  }
  size = source[0].getSize();

  graph.resize(size[0], size[1]);

#ifdef DIST_NORM
  Cluster dist;
#endif

  PointNodeMap src(graph);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      Point point;
      for (int k = 0; k < CH; ++k) {
	point[k] = source[k][i][j];
      }
      src[graph(i, j)] = point;
#ifdef DIST_NORM
      dist += point;
#endif
    }
  }
#ifdef DIST_NORM
  Matrix cm(CH, CH);
  for (int i = 0; i < CH; ++i) {
    for (int j = 0; j < CH; ++j) {
      cm(i, j) = dist.covariance(i, j);
    }
  }
  Matrix cmi = cm.inverse();
#endif
  for (Graph::EdgeIt it(graph); it != INVALID; ++it) {
#if defined(DIST_EUCLID)
    weight[it] = distance_euclid(src[graph.u(it)],
				 src[graph.v(it)]);
#elif defined(DIST_NORM)
    Vector dv(CH);
    for (int i = 0; i < CH; ++i) {
      dv(i) = src[graph.u(it)][i] - src[graph.v(it)][i];
    }
    weight[it] = dv * cmi * dv;
#endif
  }
}

template <typename Map>
struct Less {
  const Map& _map;
  Less(const Map& map) : _map(map) {}
  bool operator()(const typename Map::Key& left,
		  const typename Map::Key& right) {
    return _map[left] < _map[right];
  }
};

int main(int argc, const char *argv[]) {
  assert(argc == 3);

  Graph graph(0, 0);

  WeightEdgeMap weight(graph);
  vector<SInteger> size;

  Read(argv[1], graph, weight, size);

  IntNodeMap unionfindMap(graph, -1);
  UnionFind<IntNodeMap> unionfind(unionfindMap);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      unionfind.insert(graph(i, j));
    }
  }
  std::vector<Edge> ints(size[0] * size[1], INVALID);
  std::vector<int> sizes(size[0] * size[1], 1);

  std::vector<Edge> edges;
  for (EdgeIt it(graph); it != INVALID; ++it) {
    edges.push_back(it);
  }

  for (int i = 0; i < edges.size(); ++i) {
    std::swap(edges[i], edges[rnd[edges.size() - i] + i]);
  }

  sort(edges.begin(), edges.end(), Less<WeightEdgeMap>(weight));

  for (int i = 0; i < edges.size(); ++i) {
    Edge con = edges[i];

    int comp1 = unionfind.find(graph.u(con));
    int comp2 = unionfind.find(graph.v(con));
    if (comp1 == comp2) continue;

    Edge int1 = ints[comp1], int2 = ints[comp2];
    int size1 = sizes[comp1], size2 = sizes[comp2];
    
    if (weight[con] <= std::min(TH / size1 + weight[int1], 
				TH / size2 + weight[int2])) {
      sizes[comp1] = sizes[comp2] = size1 + size2;
      ints[comp1] = ints[comp2] = con;
      unionfind.join(graph.u(con), graph.v(con));
    }
  }

  SDataView<int> seg(size);

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      seg[i][j] = unionfind.find(graph(i, j));
    }
  }

  SDataView<unsigned char> trg(size);

  WriteSegments(seg, trg);
  

  SFile trgfile(argv[2], SFile::CREATE);
  trgfile.createDataset("segmented", sds::UINT8, size).
    setDataView<unsigned char>(trg);

  return 0;
}
	
